import keras
import os
import util
import SimpleITK as stk
import tensorflow as tf
import numpy as np
from model import unet
from keras.utils import to_categorical

path = '../build/testdata_OAR'
testresult_path = '../build/testresult'
target = ['OAR123','OAR4','OAR5','OAR6']

weights = [0,1,2,3,6,5,4]

# 参数
n_classes = 7
input_height = 128
input_width = 128

origin    = None
direction = None
space     = None

x_feed = tf.placeholder(tf.uint8,shape=[None,512,512])
y_feed = tf.placeholder(tf.uint8,shape=[None,512,512])

label_  = tf.one_hot(x_feed,n_classes,1,0)
out_    = tf.one_hot(y_feed,n_classes,1,0)
label_  = tf.cast(label_,dtype=tf.float32)
out_    = tf.cast(out_,dtype=tf.float32)

global_dice = util.global_dice(label_,out_)
dice_1 = util.dice1(label_,out_)
dice_2 = util.dice2(label_,out_)
dice_3 = util.dice3(label_,out_)
dice_4 = util.dice4(label_,out_)
dice_5 = util.dice5(label_,out_)
dice_6 = util.dice6(label_,out_)

sess = tf.Session()

for name in os.listdir(path):
    data_nii = stk.ReadImage(os.path.join('../build/testresult/{}_true_data.nii.gz'.format(name)))
    label_nii = stk.ReadImage(os.path.join('../build/testresult/{}_true_label.nii.gz'.format(name)))
    
    data = stk.GetArrayFromImage(data_nii)
    The_truth = stk.GetArrayFromImage(label_nii)

    # nii的文件元信息
    origin    = data_nii.GetOrigin()
    direction = data_nii.GetDirection()
    space     = data_nii.GetSpacing()
    
    result = []

    for i in range(len(target)):
        temp_path = os.path.join('../build/testresult/{}_{}_test_label.nii.gz'.format(name,target[i]))
        result.append(util.read_img(temp_path))
        print(result[i].shape)


    out = np.zeros(data.shape).astype(np.int16)
    for i in range(data.shape[0]):
        # 遍历每一张像素
        for j in range(data.shape[1]):
            for k in range(data.shape[2]):
                out[i,j,k] = result[0][i,j,k]
                for l in range(1,len(result)):
                    if(weights[out[i,j,k]] <= weights[result[l][i,j,k]]):
                        out[i,j,k] = result[l][i,j,k]
    

    savedImg = stk.GetImageFromArray(out)
    savedImg.SetOrigin(origin)
    savedImg.SetDirection(direction)
    savedImg.SetSpacing(space)
    stk.WriteImage(savedImg,'{}/{}_{}_test_label.nii.gz'.format(testresult_path,name,'end'))

    result = sess.run([global_dice,dice_1,dice_2,dice_3,dice_4,dice_5,dice_6],feed_dict={x_feed:The_truth,y_feed:out})
    np.save('{}/{}_{}.npy'.format(testresult_path,name,'end'),result)
    print(result)

sess.close()
